import os
import json
import subprocess
import shutil

import cv2

from utils.logger import get_logger
from PIL import Image
from config import Config

logger = get_logger(__name__)


class VideoBackgroundReplacer:
    """增强版视频背景替换器"""

    def __init__(self, task_id):
        self.task_id = task_id
        self.work_dir = os.path.join(Config.TEMP_FOLDER, f'task_{task_id}')
        os.makedirs(self.work_dir, exist_ok=True)

    def check_ffmpeg(self):
        """检查ffmpeg是否可用"""
        try:
            subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True, encoding='utf-8', errors='replace')
            return True
        except (subprocess.CalledProcessError, FileNotFoundError):
            return False

    def check_ffprobe(self):
        """检查ffprobe是否可用"""
        try:
            subprocess.run(['ffprobe', '-version'], capture_output=True, check=True, encoding='utf-8', errors='replace')
            return True
        except (subprocess.CalledProcessError, FileNotFoundError):
            return False

    def get_video_info_ffmpeg(self, video_path):
        """使用ffmpeg获取视频详细信息"""
        if not self.check_ffprobe():
            logger.error("ffprobe不可用")
            return None

        cmd = [
            'ffprobe', '-v', 'quiet', '-print_format', 'json',
            '-show_streams', '-show_format', video_path
        ]

        try:
            result = subprocess.run(cmd, capture_output=True, text=True, check=True, encoding='utf-8', errors='replace')
            info = json.loads(result.stdout)

            video_stream = None
            audio_stream = None

            for stream in info['streams']:
                if stream['codec_type'] == 'video' and video_stream is None:
                    video_stream = stream
                elif stream['codec_type'] == 'audio' and audio_stream is None:
                    audio_stream = stream

            if video_stream:
                video_info = {
                    'width': int(video_stream.get('width', 0)),
                    'height': int(video_stream.get('height', 0)),
                    'fps': eval(video_stream.get('r_frame_rate', '0/1')) if video_stream.get(
                        'r_frame_rate') != '0/0' else 25,
                    'duration': float(video_stream.get('duration', info.get('format', {}).get('duration', 0))),
                    'pixel_format': video_stream.get('pix_fmt', ''),
                    'has_alpha': 'a' in video_stream.get('pix_fmt', '').lower(),
                    'has_audio': audio_stream is not None
                }

                logger.info(f"视频信息: {video_path} - {video_info['width']}x{video_info['height']}, "
                            f"fps={video_info['fps']}, has_audio={video_info['has_audio']}")

                return video_info

        except Exception as e:
            logger.error(f"获取视频信息失败: {e}")
            return None

    def generate_thumbnail(self, video_path, output_path, time_offset=1):
        """生成视频缩略图"""
        if not self.check_ffmpeg():
            return False

        try:
            cmd = [
                'ffmpeg', '-i', video_path,
                '-ss', str(time_offset),
                '-vframes', '1',
                '-y', output_path
            ]
            subprocess.run(cmd, check=True, capture_output=True, encoding='utf-8', errors='replace')
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"生成缩略图失败: {e}")
            return False

    def generate_preview_image(self, human_path, bg_path, params, output_path):
        """生成参数预览图"""
        try:
            # 获取视频第一帧作为预览
            if self._is_video_file(human_path):
                temp_frame = os.path.join(self.work_dir, 'preview_frame.png')
                if not self._extract_frame(human_path, temp_frame):
                    return False
                human_img = Image.open(temp_frame).convert('RGBA')
            else:
                human_img = Image.open(human_path).convert('RGBA')

            # 处理背景
            if self._is_video_file(bg_path):
                temp_bg_frame = os.path.join(self.work_dir, 'preview_bg_frame.png')
                if not self._extract_frame(bg_path, temp_bg_frame):
                    return False
                bg_img = Image.open(temp_bg_frame).convert('RGB')
            else:
                bg_img = Image.open(bg_path).convert('RGB')

            # 应用参数进行预览合成
            result_img = self._compose_images(human_img, bg_img, params)

            # 保存预览图
            result_img.save(output_path)
            return True

        except Exception as e:
            logger.error(f"生成预览图失败: {e}")
            return False

    def _is_video_file(self, file_path):
        """判断是否为视频文件"""
        ext = os.path.splitext(file_path)[1].lower()
        return ext in Config.ALLOWED_VIDEO_EXTENSIONS

    def _extract_frame(self, video_path, output_path, time_offset=1.0):
        """从视频中提取帧"""
        try:
            cmd = [
                'ffmpeg', '-ss', str(time_offset),
                '-i', video_path,
                '-ss', str(time_offset),
                '-vframes', '1',
                '-y', output_path
            ]
            subprocess.run(cmd, check=True, capture_output=True, encoding='utf-8', errors='replace')
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"提取视频帧失败: {e}")
            return False

    def _compose_images(self, human_img, bg_img, params):
        """合成图像"""
        canvas_width = params.get('canvas_width', Config.DEFAULT_CANVAS_WIDTH)
        canvas_height = params.get('canvas_height', Config.DEFAULT_CANVAS_HEIGHT)

        # 添加日志
        logger.info(f"合成图像参数: canvas={canvas_width}x{canvas_height}, params={params}")

        # 处理背景
        keep_bg_ratio = params.get('keep_bg_ratio', True)
        bg_scale = params.get('bg_scale', 1.0)

        logger.info(f"背景参数: keep_ratio={keep_bg_ratio}, scale={bg_scale}, "
                    f"原始尺寸={bg_img.width}x{bg_img.height}")

        if keep_bg_ratio:
            bg_ratio = min(canvas_width / bg_img.width, canvas_height / bg_img.height)
            new_bg_w = int(bg_img.width * bg_ratio * bg_scale)
            new_bg_h = int(bg_img.height * bg_ratio * bg_scale)
        else:
            new_bg_w = int(canvas_width * bg_scale)
            new_bg_h = int(canvas_height * bg_scale)

        final_bg = Image.new('RGB', (canvas_width, canvas_height), (0, 0, 0))
        bg_img = bg_img.resize((new_bg_w, new_bg_h), Image.Resampling.LANCZOS)

        bg_x = params.get('bg_x', 0)
        bg_y = params.get('bg_y', 0)
        bg_paste_x = bg_x + (canvas_width - new_bg_w) // 2
        bg_paste_y = bg_y + (canvas_height - new_bg_h) // 2

        logger.info(f"背景位置: 缩放后尺寸={new_bg_w}x{new_bg_h}, "
                    f"粘贴位置=({bg_paste_x}, {bg_paste_y})")

        final_bg.paste(bg_img, (bg_paste_x, bg_paste_y))

        # 处理前景
        video_scale = params.get('video_scale', 1.0)
        keep_human_ratio = params.get('keep_human_ratio', True)

        logger.info(f"前景参数: keep_ratio={keep_human_ratio}, scale={video_scale}, "
                    f"原始尺寸={human_img.width}x{human_img.height}")

        if keep_human_ratio:
            human_ratio = min(canvas_width / human_img.width, canvas_height / human_img.height)
            frame_w = int(human_img.width * human_ratio * video_scale)
            frame_h = int(human_img.height * human_ratio * video_scale)
        else:
            frame_w = int(human_img.width * video_scale)
            frame_h = int(human_img.height * video_scale)

        human_img = human_img.resize((frame_w, frame_h), Image.Resampling.LANCZOS)

        # 合成
        video_x = params.get('video_x', 0)
        video_y = params.get('video_y', 0)
        paste_x = video_x + (canvas_width - frame_w) // 2
        paste_y = video_y + (canvas_height - frame_h) // 2

        logger.info(f"前景位置: 缩放后尺寸={frame_w}x{frame_h}, "
                    f"粘贴位置=({paste_x}, {paste_y})")

        final_bg.paste(human_img, (paste_x, paste_y), human_img)

        return final_bg

    def process_video_with_image_background(self, input_video_path, background_image_path,
                                            output_path, params, progress_callback=None):
        """图片背景替换"""
        if not self.check_ffmpeg():
            return False, "ffmpeg不可用"

        try:
            # 根据输出格式调整输出路径
            output_format = params.get('output_format', 'mp4').lower()
            if not output_path.endswith(f'.{output_format}'):
                output_path = output_path.rsplit('.', 1)[0] + f'.{output_format}'

            # 生成预览图作为结果缩略图
            preview_path = output_path.replace('.mp4', '_preview.jpg')
            self.generate_preview_image(input_video_path, background_image_path, params, preview_path)

            if progress_callback:
                progress_callback(25, "准备处理...")

            # 构建ffmpeg命令
            cmd = self._build_ffmpeg_command_for_image_bg(
                input_video_path, background_image_path, output_path, params
            )

            if progress_callback:
                progress_callback(50, "处理中...")

            # 执行ffmpeg命令
            result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='replace')

            if result.returncode != 0:
                logger.error(f"ffmpeg执行失败: {result.stderr}")
                return False, f"处理失败: {result.stderr}"

            if progress_callback:
                progress_callback(100, "处理完成")

            return True, "处理成功"

        except Exception as e:
            logger.error(f"处理失败: {e}")
            return False, f"处理失败: {str(e)}"

    def process_video_with_video_background(self, input_video_path, background_video_path,
                                            output_path, params, progress_callback=None):
        """视频背景替换"""
        if not self.check_ffmpeg():
            return False, "ffmpeg不可用"

        try:
            # 根据输出格式调整输出路径
            output_format = params.get('output_format', 'mp4').lower()
            if not output_path.endswith(f'.{output_format}'):
                output_path = output_path.rsplit('.', 1)[0] + f'.{output_format}'

            # 生成预览图
            preview_path = output_path.replace('.mp4', '_preview.jpg')
            self.generate_preview_image(input_video_path, background_video_path, params, preview_path)

            if progress_callback:
                progress_callback(25, "准备处理...")

            # 构建ffmpeg命令
            cmd = self._build_ffmpeg_command_for_video_bg(
                input_video_path, background_video_path, output_path, params
            )

            if progress_callback:
                progress_callback(50, "处理中...")

            # 执行ffmpeg命令
            result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='replace')

            if result.returncode != 0:
                logger.error(f"ffmpeg执行失败: {result.stderr}")
                return False, f"处理失败: {result.stderr}"

            if progress_callback:
                progress_callback(100, "处理完成")

            return True, "处理成功"

        except Exception as e:
            logger.error(f"处理失败: {e}")
            return False, f"处理失败: {str(e)}"

    def cleanup(self):
        """清理任务的临时文件"""
        try:
            if os.path.exists(self.work_dir):
                shutil.rmtree(self.work_dir)
                logger.info(f"已清理任务 {self.task_id} 的临时文件")
        except Exception as e:
            logger.error(f"清理临时文件失败: {e}")

    def _build_ffmpeg_command_for_image_bg(self, input_video, background_image, output, params):
        """构建图片背景的ffmpeg命令"""
        canvas_width = params.get('canvas_width', Config.DEFAULT_CANVAS_WIDTH)
        canvas_height = params.get('canvas_height', Config.DEFAULT_CANVAS_HEIGHT)

        # 背景参数
        bg_scale = params.get('bg_scale', 1.0)
        bg_x = params.get('bg_x', 0)
        bg_y = params.get('bg_y', 0)
        keep_bg_ratio = params.get('keep_bg_ratio', True)

        # 视频参数
        video_scale = params.get('video_scale', 1.0)
        video_x = params.get('video_x', 0)
        video_y = params.get('video_y', 0)
        keep_human_ratio = params.get('keep_human_ratio', True)

        # 获取前景视频信息（用于确定输出时长）
        video_info = self.get_video_info_ffmpeg(input_video)
        if not video_info:
            logger.warning("无法获取视频信息，使用默认设置")
            video_duration = None
        else:
            video_duration = video_info.get('duration', None)
            logger.info(f"前景视频时长: {video_duration}秒")

        # 构建复杂的 filter_complex
        filter_parts = []

        # 1. 创建黑色画布
        filter_parts.append(f'color=c=black:s={canvas_width}x{canvas_height}:d=1[canvas]')

        # 2. 处理背景图片
        if keep_bg_ratio:
            filter_parts.append(f'[0:v]scale=w={canvas_width}:h={canvas_height}:force_original_aspect_ratio=decrease,scale=w=iw*{bg_scale}:h=ih*{bg_scale}[bg_scaled]')
        else:
            bg_width = int(canvas_width * bg_scale)
            bg_height = int(canvas_height * bg_scale)
            filter_parts.append(f'[0:v]scale={bg_width}:{bg_height}[bg_scaled]')

        # 计算背景位置
        bg_overlay_x = f'(W-w)/2+{bg_x}'
        bg_overlay_y = f'(H-h)/2+{bg_y}'
        filter_parts.append(f'[canvas][bg_scaled]overlay={bg_overlay_x}:{bg_overlay_y}[bg_final]')

        # 3. 处理前景视频
        if keep_human_ratio:
            filter_parts.append(f'[1:v]scale=w={canvas_width}:h={canvas_height}:force_original_aspect_ratio=decrease,scale=w=iw*{video_scale}:h=ih*{video_scale},format=yuva420p[fg_scaled]')
        else:
            if video_info:
                fg_width = int(video_info['width'] * video_scale)
                fg_height = int(video_info['height'] * video_scale)
            else:
                fg_width = int(canvas_width * video_scale)
                fg_height = int(canvas_height * video_scale)
            filter_parts.append(f'[1:v]scale={fg_width}:{fg_height},format=yuva420p[fg_scaled]')

        # 计算前景位置
        fg_overlay_x = f'(W-w)/2+{video_x}'
        fg_overlay_y = f'(H-h)/2+{video_y}'
        filter_parts.append(f'[bg_final][fg_scaled]overlay={fg_overlay_x}:{fg_overlay_y}:format=auto')

        # 组合所有滤镜
        filter_complex = ';'.join(filter_parts)

        # 获取输出格式
        output_format = params.get('output_format', 'mp4').lower()

        # 构建基础命令
        cmd = [
            'ffmpeg',
            '-loop', '1',  # 循环背景图片
            '-i', background_image,
            '-i', input_video,
            '-filter_complex', filter_complex,
        ]

        # 添加视频编码设置
        if output_format == 'mp4':
            cmd.extend([
                '-c:v', 'libx264',
                '-crf', '23',
                '-preset', 'fast',
                '-pix_fmt', 'yuv420p',
                '-movflags', '+faststart'
            ])
        elif output_format == 'avi':
            cmd.extend([
                '-c:v', 'mpeg4',
                '-vtag', 'DIVX',
                '-qscale:v', '5'
            ])
        elif output_format == 'mov':
            cmd.extend([
                '-c:v', 'libx264',
                '-crf', '23',
                '-preset', 'fast',
                '-pix_fmt', 'yuv420p'
            ])
        elif output_format == 'webm':
            cmd.extend([
                '-c:v', 'libvpx-vp9',
                '-crf', '30',
                '-b:v', '0'
            ])

        # 设置输出时长（以前景视频为准）
        if video_duration:
            cmd.extend(['-t', str(video_duration)])
        else:
            # 如果无法获取时长，使用 shortest（但这通常不会发生在图片背景的情况下）
            cmd.extend(['-shortest'])

        # 复制前景视频的音频（如果有）
        if video_info and video_info.get('has_audio', False):
            cmd.extend([
                '-map', '1:a',  # 只映射前景视频的音频
                '-c:a', 'aac' if output_format != 'webm' else 'libopus',
                '-b:a', '128k'
            ])
        else:
            logger.info("前景视频没有音频轨道")

        cmd.extend(['-y', output])

        logger.info(f"FFmpeg命令: {' '.join(cmd)}")
        return cmd

    def _build_ffmpeg_command_for_video_bg(self, input_video, background_video, output, params):
        """构建视频背景的ffmpeg命令"""
        canvas_width = params.get('canvas_width', Config.DEFAULT_CANVAS_WIDTH)
        canvas_height = params.get('canvas_height', Config.DEFAULT_CANVAS_HEIGHT)

        # 背景参数
        bg_scale = params.get('bg_scale', 1.0)
        bg_x = params.get('bg_x', 0)
        bg_y = params.get('bg_y', 0)
        keep_bg_ratio = params.get('keep_bg_ratio', True)

        # 视频参数
        video_scale = params.get('video_scale', 1.0)
        video_x = params.get('video_x', 0)
        video_y = params.get('video_y', 0)
        keep_human_ratio = params.get('keep_human_ratio', True)

        # 获取前景和背景视频信息
        fg_info = self.get_video_info_ffmpeg(input_video)
        bg_info = self.get_video_info_ffmpeg(background_video)

        if not fg_info:
            logger.warning("无法获取前景视频信息")
            fg_duration = None
        else:
            fg_duration = fg_info.get('duration', None)
            logger.info(f"前景视频时长: {fg_duration}秒, 音频: {fg_info.get('has_audio', False)}")

        if bg_info:
            bg_duration = bg_info.get('duration', None)
            logger.info(f"背景视频时长: {bg_duration}秒")

        # 构建复杂的 filter_complex
        filter_parts = []

        # 1. 创建黑色画布
        filter_parts.append(f'color=c=black:s={canvas_width}x{canvas_height}[canvas]')

        # 2. 处理背景视频
        if keep_bg_ratio:
            filter_parts.append(f'[0:v]scale=w={canvas_width}:h={canvas_height}:force_original_aspect_ratio=decrease,scale=w=iw*{bg_scale}:h=ih*{bg_scale}[bg_scaled]')
        else:
            bg_width = int(canvas_width * bg_scale)
            bg_height = int(canvas_height * bg_scale)
            filter_parts.append(f'[0:v]scale={bg_width}:{bg_height}[bg_scaled]')

        # 计算背景位置
        bg_overlay_x = f'(W-w)/2+{bg_x}'
        bg_overlay_y = f'(H-h)/2+{bg_y}'
        filter_parts.append(f'[canvas][bg_scaled]overlay={bg_overlay_x}:{bg_overlay_y}[bg_final]')

        # 3. 处理前景视频
        if keep_human_ratio:
            filter_parts.append(f'[1:v]scale=w={canvas_width}:h={canvas_height}:force_original_aspect_ratio=decrease,scale=w=iw*{video_scale}:h=ih*{video_scale},format=yuva420p[fg_scaled]')
        else:
            if fg_info:
                fg_width = int(fg_info['width'] * video_scale)
                fg_height = int(fg_info['height'] * video_scale)
            else:
                fg_width = int(canvas_width * video_scale)
                fg_height = int(canvas_height * video_scale)
            filter_parts.append(f'[1:v]scale={fg_width}:{fg_height},format=yuva420p[fg_scaled]')

        # 计算前景位置
        fg_overlay_x = f'(W-w)/2+{video_x}'
        fg_overlay_y = f'(H-h)/2+{video_y}'
        filter_parts.append(f'[bg_final][fg_scaled]overlay={fg_overlay_x}:{fg_overlay_y}:format=auto')

        # 组合所有滤镜
        filter_complex = ';'.join(filter_parts)

        # 获取输出格式
        output_format = params.get('output_format', 'mp4').lower()

        # 构建基础命令
        cmd = ['ffmpeg']

        # 如果背景视频较短，使用 stream_loop 让它循环
        if fg_duration and bg_duration and bg_duration < fg_duration:
            loop_times = int(fg_duration / bg_duration) + 1
            cmd.extend(['-stream_loop', str(loop_times)])
            logger.info(f"背景视频将循环 {loop_times} 次")

        # 添加输入文件
        cmd.extend([
            '-i', background_video,
            '-i', input_video,
            '-filter_complex', filter_complex,
        ])

        # 添加视频编码设置
        if output_format == 'mp4':
            cmd.extend([
                '-c:v', 'libx264',
                '-crf', '23',
                '-preset', 'fast',
                '-pix_fmt', 'yuv420p',
                '-movflags', '+faststart'
            ])
        elif output_format == 'avi':
            cmd.extend([
                '-c:v', 'mpeg4',
                '-vtag', 'DIVX',
                '-qscale:v', '5'
            ])
        elif output_format == 'mov':
            cmd.extend([
                '-c:v', 'libx264',
                '-crf', '23',
                '-preset', 'fast',
                '-pix_fmt', 'yuv420p'
            ])
        elif output_format == 'webm':
            cmd.extend([
                '-c:v', 'libvpx-vp9',
                '-crf', '30',
                '-b:v', '0'
            ])

        # 设置输出时长（以前景视频为准）
        if fg_duration:
            cmd.extend(['-t', str(fg_duration)])

        # 只复制前景视频的音频
        if fg_info and fg_info.get('has_audio', False):
            cmd.extend([
                '-map', '1:a',  # 只映射前景视频（数字人）的音频
                '-c:a', 'aac' if output_format != 'webm' else 'libopus',
                '-b:a', '128k'
            ])
            logger.info("将保留前景视频的音频")
        else:
            logger.info("前景视频没有音频轨道")

        cmd.extend(['-y', output])

        logger.info(f"FFmpeg命令: {' '.join(cmd)}")
        return cmd


class VideoProcessor:
    """视频处理管理器"""

    def __init__(self):
        self.active_processors = {}

    def create_processor(self, task_id):
        """创建视频处理器"""
        processor = VideoBackgroundReplacer(task_id)
        self.active_processors[task_id] = processor
        return processor

    def get_processor(self, task_id):
        """获取视频处理器"""
        return self.active_processors.get(task_id)

    def remove_processor(self, task_id):
        """移除视频处理器"""
        if task_id in self.active_processors:
            processor = self.active_processors.pop(task_id)
            processor.cleanup()

    def cleanup_all(self):
        """清理所有处理器"""
        for task_id in list(self.active_processors.keys()):
            self.remove_processor(task_id)